#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import json
import logging
import os
import time
from multiprocessing import Process, Event
from typing import Dict, Any

import ujson

from app.services.data_identify_serv import DataIdentifyServ
from utils.u_kafka import UKafkaConsumer, UKafkaProducer
from utils.u_log import ULog


class ClassIdentifyTask(Process):
    def __init__(
            self,
            task_name,
            task_logger: logging.Logger,
            data_src_config,
            data_dst_config,
            stop_event: Event
    ):
        super().__init__(None)
        self.kafka_producer: UKafkaProducer = None
        self.kafka_consumer: UKafkaConsumer = None
        self.stop_event: Event = stop_event
        self.task_name = task_name
        if task_logger:
            self.logger = task_logger
        else:
            log_config = json.loads(os.environ['LOG_CONFIG'])
            self.logger = ULog.get_logger(self.task_name, log_config)
        self._running = False
        self.data_src_config = data_src_config
        self.data_dst_config = data_dst_config
        self.data_identify_serv = DataIdentifyServ(self.logger)

    def run(self):
        self.logger.debug(f"{self.task_name} is begin running...")
        self.logger.debug(f"stop event: {self.stop_event}")
        self.init_consumer(self.data_src_config)
        self.init_producer(self.data_dst_config)

        self.main_entry()

    def data_identify(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """执行数据分类识别"""
        try:
            _, result = self.data_identify_serv.data_identify(data)
            return result
        except Exception as e:
            self.logger.error(f"数据识别失败: {str(e)}", exc_info=True)
            return {
                'code': 500,
                'message': f"Internal error: {str(e)}",
                'data': data
            }

    def init_consumer(self, data_src_config):
        if data_src_config['source_type'] == 'kafka':
            self.kafka_consumer = UKafkaConsumer({
                'bootstrap_servers': data_src_config['bootstrap_servers'],
                'group_id': self.task_name
            }, topic=data_src_config['topic'], logger=self.logger)

    def init_producer(self, data_dst_config):
        if data_dst_config['source_type'] == 'kafka':
            self.kafka_producer = UKafkaProducer({
                'bootstrap_servers': data_dst_config['bootstrap_servers']
            }, topic=data_dst_config['topic'], logger=self.logger)

    def main_entry(self):
        while not self.stop_event.is_set():
            try:
                data = self.kafka_consumer.poll()
                if data:
                    result = self.data_identify(data)
                    result_str = ujson.dumps(result, ensure_ascii=False)
                    self.logger.debug(f"result: {result_str}")
                    self.kafka_producer.send(result_str)
            except Exception as e:
                self.logger.exception(str(e))
        self.logger.info(f"{self.task_name} is stopped!")
